import sys
from random import randint
from models.robot import Robot

# boundary x, boundary y, count for generator expression
def gen_coords(bx, by, c):
    for i in range(c):
        yield (randint(bx, by), randint(bx, by))

def dist(ab, xy):
    return abs(ab[0] + xy[0]) + abs(ab[1] + xy[1])

def main(argv):
    dust = [particle for particle in gen_coords(0, 9, 10)]
    print('dust particles', dust)

    robot = Robot(40)

    dust.sort(key=lambda d: d[0] + d[1])
    print('dust sorted   ', dust)

    while len(dust) > 0:
        clo, ind = robot.closest(dust)
        print('next closest particle would be', clo, 'at index', ind)

        # if the way to the closest particle + the way back is possible with
        # the energy that is left
        way = robot.dist(clo)
        ori = dist(clo, (0, 0))
        if way + ori <= robot.energy():
            # drive to closest
            robot.drive(clo, way)
            print(clo, 'popped')
            dust.pop(ind)
        else:
            print(way, '+', ori, 'is too far away, robot has', robot.energy(), 'energy left')
            way = robot.dist((0, 0))
            robot.drive((0, 0), way)

    print(dust)

if __name__ == "__main__":
    main(sys.argv)
